package com.ksc.wordcount.driver;

import com.ksc.wordcount.rpc.Driver.DriverRpc;
import com.ksc.wordcount.task.Task;
import com.ksc.wordcount.task.TaskContext;
import com.ksc.wordcount.task.TaskStatus;
import com.ksc.wordcount.task.TaskStatusEnum;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.BlockingQueue;

/**
 * TaskScheduler是任务调度器类，用于调度和管理任务的执行。
 * 它与TaskManager和ExecutorManager配合工作，根据任务队列和可用的执行器资源，将任务分配给合适的执行器进行处理。
 */
public class TaskScheduler {

    private TaskManager taskManager;
    private ExecutorManager executorManager;

    /**
     * 任务ID和执行器URL的映射关系
     */
    private Map<Integer, String> taskExecutorMap = new HashMap<>();

    public TaskScheduler(TaskManager taskManager, ExecutorManager executorManager) {
        this.taskManager = taskManager;
        this.executorManager = executorManager;
    }

    /**
     * 提交指定stage的任务
     * @param stageId 阶段ID
     */
    public void submitTask(int stageId) {
        // 得到stageId对应的阻塞队列，并准备取出队列中的所有任务进行任务分配
        BlockingQueue<TaskContext> taskQueue = taskManager.getBlockingQueue(stageId);

        // 循环遍历任务队列，直到队列为空
        while (!taskQueue.isEmpty()) {
            // 遍历可用执行器的资源情况
            // todo 实现 轮询给各个executor派发任务
            executorManager.getExecutorAvailableCoresMap()
                    .forEach((executorUrl, availableCores) -> {
                if (availableCores > 0 && !taskQueue.isEmpty()) {
                    // 从任务队列中取出一个任务
                    TaskContext task = taskQueue.poll();
                    // 建立任务ID和执行器URL的映射关系
                    taskExecutorMap.put(task.getTaskId(), executorUrl);
                    // 更新执行器的可用核心数 -1
                    executorManager.updateExecutorAvailableCores(executorUrl, -1);
                    // 提交任务到执行器
                    DriverRpc.submit(executorUrl, task);
                }
            });

            try {
                // 输出调试信息，并休眠1秒钟
                String executorAvailableCoresMapStr = executorManager.getExecutorAvailableCoresMap().toString();
                System.out.println("TaskScheduler submitTask stageId:" + stageId + ", taskQueue size:" + taskQueue.size()
                        + ", executorAvailableCoresMap:" + executorAvailableCoresMapStr + ", sleep 1000");
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 等待指定阶段的任务完成
     * @param stageId 阶段ID
     */
    public void waitStageFinish(int stageId) {
        // 获取阶段的任务状态
        StageStatusEnum stageStatusEnum = taskManager.getStageTaskStatus(stageId);
        // 循环等待阶段任务完成或失败
        while (stageStatusEnum == StageStatusEnum.RUNNING) {
            try {
                // 输出调试信息，并休眠1秒钟
                System.out.println("TaskScheduler waitStageFinish stageId:" + stageId + ", sleep 1000");
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            // 再次获取阶段的任务状态
            stageStatusEnum = taskManager.getStageTaskStatus(stageId);
        }

        // 检查阶段任务的最终状态，并根据需要进行处理
        if (stageStatusEnum == StageStatusEnum.FAILED) {
            System.err.println("stageId:" + stageId + " failed");
            System.exit(1);
        }
    }

    /**
     * 更新任务状态时的处理
     * @param taskStatus 任务状态
     */
    public void updateTaskStatus(TaskStatus taskStatus) {
        //任务完成或者失败则更新任务
        if (taskStatus.getTaskStatus().equals(TaskStatusEnum.FINISHED)
                || taskStatus.getTaskStatus().equals(TaskStatusEnum.FAILED)) {
            // 获取任务对应的执行器URL
            String executorUrl = taskExecutorMap.get(taskStatus.getTaskId());
            // 更新执行器的可用核心数
            executorManager.updateExecutorAvailableCores(executorUrl, 1);
        }
    }

}